package com.gome.boot.config.elasticsearch.cluster;

import com.gome.boot.common.logger.BaseLogger;
import com.gome.boot.common.util.CommonUtils;
import com.gome.boot.common.util.JacksonUtils;
import com.gome.boot.config.elasticsearch.search.ListQueryResult;
import com.gome.boot.config.elasticsearch.search.Pager;
import com.gome.boot.config.elasticsearch.search.SearchCondition;
import com.gome.boot.config.elasticsearch.tools.IndexTool;
import com.google.common.collect.Maps;
import org.elasticsearch.action.bulk.BulkItemResponse;
import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.cluster.health.ClusterHealthStatus;
import org.springframework.data.elasticsearch.ElasticsearchException;
import org.springframework.data.elasticsearch.core.DefaultEntityMapper;
import org.springframework.data.elasticsearch.core.ElasticsearchTemplate;
import org.springframework.data.elasticsearch.core.EntityMapper;
import org.springframework.data.elasticsearch.core.aggregation.AggregatedPage;
import org.springframework.data.elasticsearch.core.query.NativeSearchQueryBuilder;
import org.springframework.data.elasticsearch.core.query.SearchQuery;
import org.springframework.data.elasticsearch.core.query.StringQuery;
import org.springframework.util.Assert;

import java.util.*;

/**
 * 集群客户端扩展工具.
 *
 * @author baoxiufeng
 */
public class ClusterClient extends BaseLogger {

    /**
     * 集群可用状态
     */
    private ClusterStats stats = ClusterStats.AVAILABLE;
    /**
     * 集群监控状态
     */
    private ClusterHealthStatus healthStats = ClusterHealthStatus.GREEN;
    /**
     * 集群名称
     */
    private String name;
    /**
     * 集群客户端
     */
    private Client client;
    /**
     * Spring集群操作模板
     */
    private ElasticsearchTemplate template;
    /**
     * 实体对象映射工具
     */
    private EntityMapper entityMapper = new DefaultEntityMapper();
    /**
     * ES索引工具类集合
     */
    private Map<Class, IndexTool> indexToolMap = Maps.newHashMap();

    /**
     * 构造方法.
     *
     * @param client 集群客户端
     * @param name 集群名称
     */
    public ClusterClient(Client client, String name) {
        this(client, new ElasticsearchTemplate(client), name);
    }

    /**
     * 构造方法.
     *
     * @param client 集群客户端
     * @param template Spring集群操作模板
     * @param name 集群名称
     */
    public ClusterClient(Client client, ElasticsearchTemplate template, String name) {
        this.client = client;
        this.template = template;
        this.name = name;
    }

    /**
     * 根据条件查询指定类型的索引对象信息.
     *
     * @param searchQuery 查询条件
     * @param clazz 查询结果类
     * @param <T> 查询结果对象类型
     * @return 查询结果对象
     */
    public <T> T get(StringQuery searchQuery, Class<T> clazz) {
        return template.queryForObject(searchQuery, clazz);
    }

    /**
     * 根据条件查询指定类型的索引对象信息.
     *
     * @param condition 查询条件
     * @param clazz 查询结果类
     * @param <T> 查询结果对象类型
     * @return 查询结果对象
     */
    public <T> T get(SearchCondition condition, Class<T> clazz) {
        T entity = null;
        if (CommonUtils.isNoneBlank(condition.getRouting(), condition.getIndex(), condition.getType(), condition.getId())) {
            GetResponse response = client.prepareGet(condition.getIndex(), condition.getType(), condition.getId())
                    .setRouting(condition.getRouting()).execute().actionGet();
            try {
                entity = entityMapper.mapToObject(response.getSourceAsString(), clazz);
            } catch (Exception e) {
                logger.error("[ClusterClient.get] Failed to map source [" + response.getSourceAsString() + "] to class " + clazz.getSimpleName(), e);
            }
        } else {
            List<T> entities = getList(condition, clazz);
            entity = entities.isEmpty() ? null : entities.get(0);
        }
        return entity;
    }

    /**
     * 根据条件查询指定类型的索引对象信息列表.
     *
     * @param condition 查询条件
     * @param clazz 查询结果类
     * @param <T> 查询结果对象类型
     * @return 查询结果对象列表
     */
    public <T> List<T> getList(SearchCondition condition, Class<T> clazz) {
        Assert.notNull(condition.getQueryBuilder(), "Cannot query by 'null' builder.");
        // 构建查询语句
        NativeSearchQueryBuilder queryBuilder = new NativeSearchQueryBuilder()
                .withQuery(condition.getQueryBuilder()).withPageable(condition.getPageable());
        if (CommonUtils.isNotEmpty(condition.getIndex())) {
            queryBuilder.withIndices(condition.getIndex());
        }
        if (CommonUtils.isNotEmpty(condition.getType())) {
            queryBuilder.withTypes(condition.getType());
        }
        if (CommonUtils.isNotEmpty(condition.getRouting())) {
            queryBuilder.withRoute(condition.getRouting());
        }
        // 执行查询
        try {
            Iterator<T> iterator = template.queryForPage(queryBuilder.build(), clazz).iterator();
            List<T> entities = new ArrayList<>();
            while (iterator.hasNext()){
                entities.add(iterator.next());
            }
            return entities;
        } catch (Exception e) {
            logger.error("[ClusterClient.getList] search failed by condition: " + JacksonUtils.convertObjectToJson(condition), e);
            return Collections.emptyList();
        }
    }

    /**
     * 根据条件查询指定类型的索引对象信息列表及分页信息.
     *
     * @param condition 查询条件
     * @param clazz 查询结果类
     * @param <T> 查询结果对象类型
     * @return 查询结果对象列表及分页信息
     */
    public <T> ListQueryResult<T> getListResult(SearchCondition condition, Class<T> clazz) {
        Assert.notNull(condition.getQueryBuilder(), "Cannot query by 'null' builder.");
        // 构建查询语句
        SearchQuery query = new NativeSearchQueryBuilder()
                .withQuery(condition.getQueryBuilder())
                .withPageable(condition.getPageable())
                .withRoute(condition.getRouting())
                .build();
        // 执行查询
        try {
            AggregatedPage<T> aggregatedPage = template.queryForPage(query, clazz);
            Pager pager = new Pager(condition.getPageable().getPageNumber() + 1, condition.getPageable().getPageSize());
            pager.setTotalPage(aggregatedPage.getTotalPages());
            pager.setTotalSize((int) aggregatedPage.getTotalElements());
            Iterator<T> iterator = aggregatedPage.iterator();
            List<T> orders = new ArrayList<>();
            while (iterator.hasNext()){
                orders.add(iterator.next());
            }
            return new ListQueryResult<T>(pager, orders);
        } catch (Exception e) {
            logger.error("[ClusterClient.getListResult] search failed by condition: " + JacksonUtils.convertObjectToJson(condition), e);
            return null;
        }
    }

    /**
     * 根据条件查询指定类型的索引对象信息数量.
     *
     * @param condition 查询条件
     * @param clazz 查询结果类
     * @param <T> 查询结果对象类型
     * @return 查询结果对象列表
     */
    public <T> Long getCount(SearchCondition condition, Class<T> clazz) {
        Assert.notNull(condition.getQueryBuilder(), "Cannot query by 'null' builder.");
        // 构建查询语句
        SearchQuery query = new NativeSearchQueryBuilder()
                .withQuery(condition.getQueryBuilder())
                .withRoute(condition.getRouting())
                .build();
        // 执行查询
        try {
            return template.count(query, clazz);
        } catch (Exception e) {
            logger.error("Search doc count failed by condition: " + JacksonUtils.convertObjectToJson(condition), e);
        }
        return 0L;
    }

    /**
     * 保存/更新索引对象信息请求.
     *
     * @param request 索引对象保存/更新请求
     * @return 保存/更新结果
     */
    public boolean save(IndexRequest request) {
        Assert.notNull(request, "Cannot save by 'null' indexRequest.");
        if (logger.isInfoEnabled()) {
            logger.info("Save index request:{}", JacksonUtils.convertObjectToJson(request));
        }
        IndexResponse indexResponse = client.index(request).actionGet();
        if (logger.isInfoEnabled()) {
            logger.info("Save index response:{}", JacksonUtils.convertObjectToJson(indexResponse));
        }
        return indexResponse.getShardInfo().getFailed() == 0;
    }

    /**
     * 删除索引对象信息请求.
     *
     * @param request 索引对象删除请求
     * @return 删除结果
     */
    public boolean delete(DeleteRequest request) {
        Assert.notNull(request, "Cannot delete by 'null' deleteRequest.");
        if (logger.isInfoEnabled()) {
            logger.info("Delete request:{}", JacksonUtils.convertObjectToJson(request));
        }
        DeleteResponse deleteResponse = client.delete(request).actionGet();
        if (logger.isInfoEnabled()) {
            logger.info("Delete response:{}", JacksonUtils.convertObjectToJson(deleteResponse));
        }
        return deleteResponse.getShardInfo().getFailed() == 0;
    }

    /**
     * 批量保存/更新索引对应信息.
     *
     * @param requestList 保存/更新索引对象请求列表
     * @return 批量保存/更新结果
     */
    public boolean bulk(String clusterName, List<IndexRequest> requestList) {
        if (CommonUtils.isEmpty(requestList)) {
            return false;
        }
        BulkRequestBuilder builder = client.prepareBulk();
        for (IndexRequest request : requestList) {
            builder.add(request);
        }
        BulkResponse response = null;
        try {
            response = builder.execute().actionGet();
        } catch (Exception e) {
            for (IndexRequest request : requestList) {
                MultiClusterLogger.expLog(clusterName, request.index(), request.type(), request.routing(), request.id());
            }
            throw e;
        }
        if (response.hasFailures()) {
            Map<String, String> failedDocuments = new HashMap<String, String>();
            for (BulkItemResponse item : response.getItems()) {
                if (item.isFailed()) {
                    failedDocuments.put(item.getId(), item.getFailureMessage());
                    MultiClusterLogger.norLog(clusterName, item.getIndex(), item.getType(), null, item.getId());
                }
            }
            throw new ElasticsearchException(
                    "Bulk indexing has failures. Use ElasticsearchException.getFailedDocuments() for detailed messages ["
                            + failedDocuments + "]", failedDocuments
            );
        }
        return true;
    }

    public void addIndexTool(Class clazz, IndexTool indexTool) {
        indexToolMap.put(clazz, indexTool);
    }

    public String getName() {
        return name;
    }

    public void setName(String name) {
        this.name = name;
    }

    public ClusterStats getStats() {
        return stats;
    }

    public void setStats(ClusterStats stats) {
        this.stats = stats;
    }

    public ClusterHealthStatus getHealthStats() {
        return healthStats;
    }

    public void setHealthStats(ClusterHealthStatus healthStats) {
        this.healthStats = healthStats;
    }

    public Client getClient() {
        return client;
    }

    public ElasticsearchTemplate getTemplate() {
        return template;
    }
}
